package com.bsj.media.media.packageFlv;

import cn.hutool.core.util.StrUtil;

import com.bsj.media.media.NalInfo;
import com.bsj.media.media.ParserH264;
import com.bsj.media.media.ParserH265;
import com.bsj.media.media.ParserVideoData;
import com.bsj.media.media.h26X.H264NaluSPSInfo;
import com.bsj.media.media.h26X.H264SliceHead;
import com.bsj.media.media.until.AACHandlerUtil;
import com.bsj.power.common.def.protocol.rtp.RtpParameterSet;
import com.bsj.tools.Conv;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;


import java.nio.ByteBuffer;
import java.util.*;

/**
 * @Description flv各打包类型处理
 * @Date 2022/10/25 14:09
 **/
@Slf4j
public abstract class FlvTypeHandler {



    public boolean isPacketFlvHeader = true; // 是否打包FLV头

    /**
     * flv头的内容标记 0x05:音视频，0x01:视频，0x04:音频
     */
    public byte flvFlags;
    /**
     * 音频相关参数：
     * audioType：        音频类型，0：ADPCM 1：G711A 2：G726 3:AAC，  默认是G711a
     * st_audioRate:      每帧音频的持续时长
     * audioSampleRate:   音频采样率
     * sampleSize:        音频采样大小
     * audiocodecid:      音频编码类型  10 ACC
     * audiotimestamps:   音频相对时间戳
     */
    public int audioType = 1;
    public int st_audioRate = 23;
    public int audioSampleRate;
    public int sampleSize = 16;
    public int audiocodecid = 10;
    public int audiotimestamps;
    /**
     * 视频相关参数：
     * st_videoRate：        每帧视频的持续时长
     * width:                视频宽
     * height:               视频高
     * ratioEnum             FD分辨率枚举
     * videoFrameRate:       帧率
     * videocodecid:         视频编码类型  7：H264  12:H265
     * videotimestamps       视频相对时间戳
     * videoFrameTimeStamps  用于存储FD上面的视频时间戳（秒）
     */
    public int st_videoRate = 23;
    public int width;
    public int height;
    public int ratioEnum;
    public int videoFrameRate;
    public byte videocodecid;
    public int videotimestamps;
    public long videoFrameTimeStamps;
    /**
     * 音视频打包转视频打包，为了解决前端请求音视频，设备确没有音频上来，导致页面卡的问题
     */
    public volatile boolean mediaToVideoPacket;
    /**
     * aac打包工具
     */
    public AACHandlerUtil aacHandlerUtil;
    /**
     * h264的sps帧详细信息，用于去解析非关键帧和关键帧的SliceHead(片段头)
     */
    H264NaluSPSInfo h264NaluSPSInfo;
    /**
     * 存放采样或未采样aac数据（去帧头）
     * 存放格式： 2字节长度+aac数据+2字节长度+aac数据。。。。。
     * 定义该变量目的： 不用频繁创建byte[]来存放aac数据，通过ByteBuffer下标来有效利用堆内存
     */
    public ByteBuffer aacDatas;
    /**
     * 存放上一个的tagSize
     */
    public int previousTagSize = 0;



    public FlvTypeHandler(byte flvFlags) {
        this.flvFlags = flvFlags;

    }

    /**
     * 第一个打包flv(flv Header + TagSize[0] + ScriptTag + 音视频第一个的Tag)
     *
     * @param fdBody: I帧数据
     * @param dst:
     * @return
     **/
    public void firstPackageFlv(byte[] fdBody, ByteBuffer dst) {
        if (preJudgement(fdBody)) {
            // 1、打包flvHeader + TagSize[0]
            FlvPacketUtil.getFlvHeader(dst, flvFlags);
            // 2、打包ScriptTag
            previousTagSize = FlvPacketUtil.setScriptTag(dst, videoFrameRate, width, height, videocodecid, audioSampleRate, sampleSize, audiocodecid, previousTagSize);
            // 3、打包音频或者视频的第一个Tag
            setFristMediaTag(fdBody, dst);
            audiotimestamps += st_audioRate;
            videotimestamps += st_videoRate;
        }
    }

    /**
     * 前置判断，成立才继续打包，获取相应的参数（音频、视频）配置
     *
     * @param fdBody:
     * @return
     **/
    public abstract boolean preJudgement(byte[] fdBody);

    /**
     * 对flv的音视频第一Tag进行打包
     *
     * @param fdBody:
     * @param dst:
     * @return
     **/
    public abstract void setFristMediaTag(byte[] fdBody, ByteBuffer dst);

    /**
     * 返回是否可以打包视频
     *
     * @return
     **/
    public abstract boolean packVideoFlag();

    /**
     * 返回是否可以打包音频
     *
     * @return
     **/
    public abstract boolean packAudioFlag();

    /**
     * 分辨率， 详见协议
     *
     * @param solution
     */
    public void getSolutionRatio(int solution) {
        ratioEnum = solution;
        switch (solution) {
            case 0:
                width = 176;
                height = 144;
                break;
            case 2:
                width = 704;
                height = 288;
                break;
            case 3:
                width = 704;
                height = 576;
                break;
            case 4:
                width = 1280;
                height = 720;
                break;
            case 5:
                width = 1280;
                height = 960;
                break;
            case 6:
                width = 1920;
                height = 1080;
                break;
            case 7:
                width = 2560;
                height = 1440;
                break;
            case 127:
                width = 960;
                height = 540;
                break;
            default:
                width = 352;
                height = 288;
                break;
        }
    }

    /**
     * 对H264的第一个tag进行打包
     *
     * @param fdBody:
     * @param dst:
     * @return
     **/
    void h264FristTagFlv(byte[] fdBody, ByteBuffer dst) {
        // 获取SPS和PPS的帧
        NalInfo sps = new NalInfo();
        NalInfo pps = new NalInfo();

        int ret = ParserH264.ReadOneNaluFromBuf(fdBody, fdBody.length, 21, sps);
        if (ret < 0 || sps.getNal_unit_type() != 0x7) {
            log.error(", 获取H264的SPS帧异常:");
            dst.position(0);
            return;
        }
        h264NaluSPSInfo = H264NaluSPSInfo.analysisNaluSPSInfo(Arrays.copyOfRange(sps.getBuf(), sps.getBuf()[2] == 0x01 ? 3 : 4, sps.getBuf().length));
        ret = ParserH264.ReadOneNaluFromBuf(fdBody, fdBody.length, ret, pps);
        if (ret < 0 || pps.getNal_unit_type() != 0x08) {
            log.error(", 获取H264的pps帧异常:");
            dst.position(0);
            return;
        }
        previousTagSize = FlvPacketUtil.setH264FirstVideoTag(dst, sps.getBuf(), pps.getBuf(), previousTagSize);
    }

    /**
     * 对aac的第一Tag进行打包
     *
     * @param dst:
     * @return
     **/
    public void aacAudioFristTagFlv(ByteBuffer dst) {
        int aacSampleRate = 3; // Flv标记：采样率44K
        int aacChannel = 0; // 立体声
        if (audioSampleRate == 16000) {
            aacSampleRate = 1;
        }
        previousTagSize = FlvPacketUtil.setAACFirstTag(dst, aacSampleRate, aacChannel, audiocodecid, previousTagSize);
    }

    /**
     * 对H265的第一个tag进行打包
     *
     * @param fdBody:
     * @param dst:
     * @return
     **/
    void h265FristTagFlv(byte[] fdBody, ByteBuffer dst) {
        NalInfo vps = new NalInfo();
        NalInfo sps = new NalInfo();
        NalInfo pps = new NalInfo();

        // 读取一帧数据 vps sps pps
        int ret = ParserH265.ReadOneNaluFromBuf(fdBody, fdBody.length, 21, vps);
        if (ret < 0 || vps.getNal_unit_type() != RtpParameterSet.naluTypeVPSH5) {
            log.error(",获取H265的vps帧异常:");
            dst.position(0);
            return;
        }
        ret = ParserH265.ReadOneNaluFromBuf(fdBody, fdBody.length, ret, sps);
        if (ret < 0 || sps.getNal_unit_type() != RtpParameterSet.naluTypeSPSH5) {
            log.error(",获取H265的sps帧异常:");
            dst.position(0);
            return;
        }
        ret = ParserH265.ReadOneNaluFromBuf(fdBody, fdBody.length, ret, pps);
        if (ret < 0 || pps.getNal_unit_type() != RtpParameterSet.naluTypePPSH5) {
            log.error(",pps:");
            dst.position(0);
            return;
        }
        previousTagSize = FlvPacketUtil.setH265FirstVideoTag(dst, vps.getBuf(), sps.getBuf(), pps.getBuf(), previousTagSize);
    }

    /**
     * 打包flv视频的tag + tagSize
     *
     * @param body:
     * @param dst:
     * @return
     **/
    public void packVideo(byte[] body, ByteBuffer dst) {
        videoFrameTimeStamps = Conv.getLongNetOrder(body, 5) / 1000;
        boolean isIDR = ParserVideoData.checkIframe(body);
        int newRatioEnum = ParserVideoData.getVideoSolution(body, 0);
        if (isIDR && newRatioEnum != ratioEnum) {
            // 分辨率改变, 重新打包flv
            log.info(StrUtil.format("{}分辨率改变， oldEnum={}, newEnum={}", "123", ratioEnum, newRatioEnum));
            ratioEnum = newRatioEnum;
            getSolutionRatio(ratioEnum);
            // 打包ScriptTag
            previousTagSize = FlvPacketUtil.setScriptTag(dst, videoFrameRate, width, height, videocodecid, audioSampleRate, sampleSize, audiocodecid, previousTagSize);
            // 3、打包音频或者视频的第一个Tag
            setFristMediaTag(body, dst);
        }
        // 目前body中，PB帧可能有多个，所以分开处理
//        if (ParserVideoData.checkIframe(body)) {
//            // I帧处理
//            packIVidoTag(body, dst);
//        } else {
//            // 其他帧处理
//            packNotIDRTag(body, dst);
//        }

        packVideoTag(body, dst, isIDR);
    }

    /**
     * 具有把相同分片合并到同一个videoTag的功能
     *
     * @param body:
     * @param dst:
     * @param isIDR:
     * @return
     **/
    public void packVideoTag(byte[] body, ByteBuffer dst, boolean isIDR) {
        int encodingFormat = ParserVideoData.getAudioType(body); // 视频格式
        int offset = isIDR ? FlvPacketUtil.findIDROffset2(body, 21, encodingFormat) : ParserH264.FindStartCodeEx(body, 21, body.length);
        if (offset <= 0) {
            log.warn(",视频帧打包失败, isIDR=" + isIDR);
            return;
        }

        // 1、对body提取全部的帧数据(去除头), key：帧编号
        Map<Integer, List<Pair<Integer, Integer>>> nalus = findVideoFrameSlice(body, offset, encodingFormat);
        // 2、遍历nalus
        nalus.forEach((k, v) -> {
            // 打包videoTag
            previousTagSize = FlvPacketUtil.setVideoTag(body, dst, videocodecid, isIDR ? 1 : 2, videotimestamps, videoFrameRate, previousTagSize, v);
            videotimestamps += st_videoRate;
        });
    }

    /**
     * 查询视频帧的分片
     *
     * @param body:
     * @param offset:
     * @param encodingFormat: 帧类型，只接受两种 1：h264 3:h265
     * @return
     **/
    public Map<Integer, List<Pair<Integer, Integer>>> findVideoFrameSlice(byte[] body, int offset, int encodingFormat) {
        Map<Integer, List<Pair<Integer, Integer>>> map = new HashMap<>();
        int frameNum = 0; // 帧的分片号，h265先不区分是否有相同的分片，都当成不一样的分片处理
        while (true) {
            // 查询开始头，再查询下一个非关键帧
            int startHeader = body[offset + 2] == 0x01 ? offset + 3 : offset + 4;
            int ret2 = ParserH264.FindStartCodeEx(body, startHeader, body.length);
            if (ret2 == -1) {
                // 说明只有一帧
                int frameLen = body.length - startHeader;
                frameNum = (encodingFormat == RtpParameterSet.fdVideoCodeTypeH265) ? frameNum++ : H264SliceHead.getFrameNum(Arrays.copyOfRange(body, startHeader, startHeader + frameLen), h264NaluSPSInfo);
                if (frameNum != -1) {
                    List<Pair<Integer, Integer>> pairs = map.computeIfAbsent(frameNum, k -> new ArrayList<>());
                    pairs.add(new ImmutablePair<>(startHeader, frameLen));
                }
                break;
            } else {
                // 说明后面还有帧
                int frameLen = ret2 - startHeader;
                frameNum = (encodingFormat == RtpParameterSet.fdVideoCodeTypeH265) ? frameNum++ : H264SliceHead.getFrameNum(Arrays.copyOfRange(body, startHeader, startHeader + frameLen), h264NaluSPSInfo);
                if (frameNum != -1) {
                    List<Pair<Integer, Integer>> pairs = map.computeIfAbsent(frameNum, k -> new ArrayList<>());
                    pairs.add(new ImmutablePair<>(startHeader, frameLen));
                }
                offset = ret2;
            }
        }
        return map;
    }

    /**
     * 打包非关键帧Tag, 主要逻辑就是body里面可能有多个非关键帧
     *
     * @param body:
     * @param dst:
     * @return
     **/
    protected void packNotIDRTag(byte[] body, ByteBuffer dst) {
        int srcOffset = 21; // 视频帧的初始偏移量
        int ret = ParserH264.FindStartCodeEx(body, srcOffset, body.length);
        if (ret > 0) {
            while (true) {
                // 查询开始头，再查询下一个非关键帧
                int startHeader = body[ret + 2] == 0x01 ? ret + 3 : ret + 4;
                int ret2 = ParserH264.FindStartCodeEx(body, startHeader, body.length);
                if (ret2 == -1) {
                    // 说明就只有一帧
                    int frameLen = body.length - startHeader;
                    previousTagSize = FlvPacketUtil.setVideoTag(body, startHeader, frameLen, dst, videocodecid, 2, videotimestamps, videoFrameRate, previousTagSize);
                    videotimestamps += st_videoRate;
                    break;
                } else if (ret2 > ret) {
                    // 说明后面还是帧, 先打包前面一帧
                    int frameLen = ret2 - startHeader;
                    previousTagSize = FlvPacketUtil.setVideoTag(body, startHeader, frameLen, dst, videocodecid, 2, videotimestamps, videoFrameRate, previousTagSize);
                    videotimestamps += st_videoRate;
                    // 设置下次查找h264 起始位置
                    ret = ret2;
                }
            }
        } else {
            log.error(",对flv非关键视频帧打包失败！！");
        }
    }

    protected void packIVidoTag(byte[] body, ByteBuffer dst) {
        // 查找关键帧的偏移量
        int encodingFormat = ParserVideoData.getAudioType(body); // 视频格式
        int idrOffset = FlvPacketUtil.findIDROffset(body, 21, encodingFormat);
        if (idrOffset != 0) {
            int frameLen = body.length - idrOffset; // 需要复制的帧大小
            previousTagSize = FlvPacketUtil.setVideoTag(body, idrOffset, frameLen, dst, videocodecid, 1, videotimestamps, videoFrameRate, previousTagSize);
            videotimestamps += st_videoRate;
        } else {
            log.error( ",对flv关键视频帧打包失败！！");
        }
    }

    /**
     * 打包flv音频的tag + tagSize
     *
     * @param body:
     * @param dst:
     * @return
     **/
    public void packAudio(byte[] body, ByteBuffer dst) {
        int audioType = ParserVideoData.getAudioType(body);
        if (aacHandlerUtil == null) {
            // 初始化aac处理
            aacDatas = ByteBuffer.allocate(5 * 1024);
            aacHandlerUtil = new AACHandlerUtil(audioType, "");
        }
        aacHandlerUtil.handlerAudioData(audioType, aacDatas, body, 21);
        aacDatas.flip(); // 切换读模式
        if (aacDatas.remaining() <= 0) {
//            log.error("重采样失败, uri:" );
            return;
        }
        // 解析多个aac数据
        while (aacDatas.remaining() > 0) {
            short aacLen = aacDatas.getShort();
            if (aacLen <= aacDatas.remaining() && aacLen <= 1024 * 10) {
                // 打包flv音频tag
                byte[] aac = new byte[aacLen];
                aacDatas.get(aac);
                previousTagSize = FlvPacketUtil.setAudioTag(aac, dst, audiocodecid, audioSampleRate == 16000 ? 1 : 3, 0, audiotimestamps, previousTagSize);
                audiotimestamps += st_audioRate;
            } else {
                log.error("解析aac数据异常: aacLen:" + aacLen + ", 缓存剩余可读长度:" + aacDatas.remaining());
                break;
            }
        }
    }

}
